# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import json

import numpy as np

from experimental.prophet.graph import graph_def_pb2
from experimental.prophet.graph.shape import Shape
from experimental.prophet.graph.util import from_data_type_proto


class PyInputVariable:
  def __init__(self, native_variable):
    self._var = native_variable
    self._arg_name = self._var.arg_name
    self._dtype = from_data_type_proto(self._var.dtype)
    self._value_ref = None

  @property
  def arg_name(self):
    return self._arg_name

  @property
  def dtype(self):
    return self._dtype

  @property
  def shape(self):
    return self._var.shape

  @property
  def is_null(self):
    return self._var.is_null()

  @property
  def is_scalar(self):
    return len(self._var.shape) == 0

  @property
  def value(self):
    if self._var.is_null():
      return None

    if self._value_ref is None or self._var.is_buffer_updated():
      self._value_ref = np.array(self._var, copy=False)

    if self.is_scalar:
      return self._value_ref.item()  # Scalar
    else:
      return self._value_ref


class PyOutputVariable(PyInputVariable):
  def __init__(self, native_variable):
    super().__init__(native_variable)

  @property
  def value(self):
    return super().value

  @value.setter
  def value(self, value):
    if value is None:
      self._var.set_null(True)
    else:
      self._var.set_value(value)


class PyOpKernelContext:
  def __init__(self, kernel_cls):
    self._kernel_cls = kernel_cls
    self._inputs = None
    self._outputs = None
    self._node_def = None
    self._attrs = None
    self._kernel = None

  def init(self, serialized_node_def, inputs, outputs):
    self._inputs = tuple([PyInputVariable(input) for input in inputs])
    self._outputs = tuple([PyOutputVariable(output) for output in outputs])

    self._node_def = graph_def_pb2.NodeDefProto()
    self._node_def.ParseFromString(serialized_node_def)

    self._attrs = None
    if self._node_def.attrs_json:
      self._attrs = json.loads(self._node_def.attrs_json)
    self._kernel = self._kernel_cls(self._attrs, self._node_def)

  def compute(self):
    self._kernel.compute(self._inputs, self._outputs)

  def on_null_input(self):
    self._kernel.on_null_input(self._inputs, self._outputs)
